import tushare as ts
from openpyxl import load_workbook
import pandas as pd
from datetime import datetime, timedelta

# 初始化Tushare接口（需替换为你的token）
ts.set_token("e0b9a82aa635937e1b86f58f025d5720b68f4e1b9c0884bb8787ad1f")
pro = ts.pro_api()


def get_trade_dates():
    """获取交易日历"""
    df = pro.trade_cal(exchange='', start_date='20200101', end_date='20251231')
    trade_dates = df[df['is_open'] == 1]['cal_date'].tolist()
    return trade_dates


def get_yesterday():
    """获取上一个交易日"""
    trade_dates = get_trade_dates()
    today = datetime.today().strftime('%Y%m%d')
    if today in trade_dates:
        idx = trade_dates.index(today)
        return trade_dates[idx - 1]
    else:
        return trade_dates[-1]


def get_financial_reports():
    """获取近5年年报数据"""
    # 生成近5年年报报告期
    end_year = datetime.now().year - 1
    report_periods = [f"{year}1231" for year in range(end_year - 4, end_year + 1)]

    # 存储所有公司财报数据
    all_reports = []

    # 获取全量股票列表
    stocks = pro.stock_basic(exchange='', list_status='L')['ts_code'].tolist()

    for stock in stocks:
        for period in report_periods:
            try:
                df = pro.fina_indicator(ts_code=stock, period=period, limit=1)
                if not df.empty:
                    df['ts_code'] = stock
                    df['report_period'] = period
                    all_reports.append(df)
            except Exception as e:
                print(f"获取{stock} {period}财报失败: {str(e)}")

    return pd.concat(all_reports)


def get_market_data():
    """获取昨日行情数据"""
    yesterday = get_yesterday()
    df = pro.daily(trade_date=yesterday)
    return df[['ts_code', 'open', 'high', 'low', 'close', 'vol']]


def write_to_excel(fin_data, market_data, template_path, output_path):
    """写入指定Excel模板"""
    # 加载模板
    wb = load_workbook(template_path)

    # 写入财报数据
    fin_sheet = wb['财报数据']
    start_row = 2  # 假设数据从第2行开始
    for row_num, row in fin_data.iterrows():
        for col_num, value in enumerate(row):
            fin_sheet.cell(row=start_row + row_num, column=col_num + 1).value = value

    # 写入行情数据
    market_sheet = wb['行情数据']
    start_row = 2
    for row_num, row in market_data.iterrows():
        for col_num, value in enumerate(row):
            market_sheet.cell(row=start_row + row_num, column=col_num + 1).value = value

    # 保存结果
    wb.save(output_path)
    print(f"数据已写入{output_path}")


if __name__ == "__main__":
    # 执行数据抓取
    financial_data = get_financial_reports()
    market_data = get_market_data()

    # 写入Excel模板（需提前准备模板文件）
    write_to_excel(
        financial_data,
        market_data,
        template_path="财报模板.xlsx",
        output_path="股市数据输出.xlsx"
    )